# -*- coding: utf-8 -*-
import logging
import time

import paramiko

logger = logging.getLogger(__name__)  # 设置日志初始模块
logger.setLevel(level=logging.INFO)  # 设置日志级别
handler = logging.FileHandler("server_commands.log")  # 设置当前目录存放日志
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(lineno)d - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)

class SSHClient(object):
    def __init__(self, hostname, username, port, passwd, **kwargs):
        self.hostname = hostname
        self.username = username
        self.port = port
        self.passwd = passwd
        self.ssh_client = None
        self.transport = None
        self.channel = None
        self.sftp = None

    def ssh_channel(self, connect_type=True):
        ssh_client = paramiko.SSHClient()
        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy)
        if connect_type == 1:
            try:
                ssh_client.connect(username=self.username, password=self.passwd, port=self.port, hostname=self.hostname)
            except Exception as e:
                logger.error(e)
                logger.error("Please check the username or password")
                return "请检查用户名密码"
        else:
            try:
                pkey = "/root/.ssh/id_rsa"
                key = paramiko.RSAKey.from_private_key(pkey)
                ssh_client.connect(hostname=self.hostname, pkey=key)
            except Exception as e:
                logger.error(e)
                logger.error("Please check the correctness of the key")
                return "请检查密钥是否正确"
        self.channel = ssh_client.get_transport().open_session()
        self.channel.get_pty()
        self.channel.invoke_shell()
        return True

    def sftp_channel(self, connect_type=True):
        try:
            self.transport = paramiko.Transport(self.hostname, self.port)
            if connect_type:
                self.transport.connect(username=self.username, password=self.passwd)
            else:
                pkey = "/root/.ssh/id_rsa"
                key = paramiko.RSAKey.from_private_key(pkey)
                self.transport.connect(username="root", pkey=key)
            self.sftp = paramiko.SFTPClient.from_transport(self.transport)
        except Exception as e:
            logger.error(e)
            logger.error("Please check the correctness of user name, password or key")
            return "请检查用户名密码或者密钥正确性"

    def close_ssh(self):
        if self.ssh_client is not None:
            self.channel.close()
            self.ssh_client.close()
        if self.sftp is not None:
            self.transport.close()
            self.sftp.close()

    def check_memory(self):
        cmd = "df -h"
        res = self.exec_command(cmd)
        return res

    def exec_command(self, command, sleep_time=2) -> str:
        # 向服务器发送命令
        try:
            self.channel.send(command + "\n")
        except Exception as e:
            logger.error("please check server connect")
            return "check server connect"
        return_data = b""
        while True:
            if self.channel.recv_ready():
                output = self.channel.recv(1024)
                return_data = return_data + output
            else:
                time.sleep(sleep_time)
                if not (self.channel.recv_ready()):
                    break
        # print(return_data.decode())
        return return_data.decode()

    def docker_image(self):
        cmd = "docker images"
        res = self.exec_command(cmd)
        print(res)


if __name__ == '__main__':
    ssh = SSHClient("10.239.183.65", "root", 22, "password")
    channel = ssh.ssh_channel()
    if channel:
        res = ssh.check_memory()
    ssh.close_ssh()






